#!/usr/bin/env python3
"""
单线程读取 ROS2 bag，按 bunch_size 帧为一步，matplotlib 3D 动画播放
python3.12 bag2mpl3d_single.py \
  --bag /gitlab/pcd_analysis/demo_data/unilidar-2023-09-22-12-42-04.db3 \
  --topic /velodyne_points \
  --bunch 1000 \
  --interval 50
"""
import argparse
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from rosbag2_py import SequentialReader, StorageOptions, ConverterOptions, StorageFilter
from sensor_msgs.msg import PointCloud2
from rclpy.serialization import deserialize_message

def pc2_to_xyz(cloud: PointCloud2):
    """把 PointCloud2 解析成 (N,3) 数组，忽略无效点"""
    a = np.frombuffer(cloud.data, dtype=np.uint8).reshape(-1, cloud.point_step)
    xyz = a[:, :12].copy().view(np.float32).reshape(-1, 3)
    return xyz[np.isfinite(xyz).all(axis=1)]

def load_bunches(bag_path: str, topic: str, bunch_size: int):
    """返回 list[ (N,3) ]，每元素是 bunch_size 帧合并后的点云"""
    reader = SequentialReader()
    reader.open(StorageOptions(uri=bag_path, storage_id='sqlite3'),
                ConverterOptions('', ''))
    reader.set_filter(StorageFilter(topics=[topic]))

    bunches, cache = [], []
    while reader.has_next():
        _, data, _ = reader.read_next()
        msg = deserialize_message(data, PointCloud2)
        cache.append(pc2_to_xyz(msg))
        if len(cache) == bunch_size:
            bunches.append(np.vstack(cache))   # 合并成一步
            cache = []
    if cache:          # 尾部不足 bunch_size
        bunches.append(np.vstack(cache))
    return bunches

def animate(i):
    """FuncAnimation 回调：更新散点"""
    pts = bunches[i]
    print(f'step {i+1}/{len(bunches)}  points={len(pts)}')
    scat._offsets3d = (pts[:, 0], pts[:, 1], pts[:, 2])
    ax.set_title(f'step {i+1}/{len(bunches)}  points={len(pts)}')
    return scat,

# ---------- main ----------
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--bag', default='/gitlab/pcd_analysis/demo_data/unilidar-2023-09-22-12-42-04.db3')
    parser.add_argument('--topic', default='/unilidar/cloud')
    parser.add_argument('--bunch', type=int, default=10, help='每步合并多少帧')
    parser.add_argument('--interval', type=int, default=50, help='毫秒/步')
    args = parser.parse_args()

    print('读取 bag 并打包……')
    global bunches
    bunches = load_bunches(args.bag, args.topic, args.bunch)
    print(f'共 {len(bunches)} 步，准备播放')

    fig = plt.figure()
    global ax, scat
    ax = fig.add_subplot(111, projection='3d')
    scat = ax.scatter([], [], [], s=1)
    ax.set_box_aspect([1,1,1])

    ani = animation.FuncAnimation(fig, animate, frames=len(bunches),
                                  interval=args.interval, blit=True, repeat=False)
    plt.show(block=True)

if __name__ == '__main__':
    main()